// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file
\brief	Main Body of the FAST_PEPI_BT_V2 MPI Program
*/
/*! 
\defgroup	fast_pepi_bt_v2 FAST_PEPI_BT_V2 - Fast Epistasis for Binary Trait
\ingroup    mpi_stats
\brief		FAST_PEPI_BT_V2 - Fast Epistasis for Binary Trait
\details
<p>
  This program is based on the fast epistasis (BT) %option of PLINK.
  It performs the same interaction calculation and 
  unlike PLINK can also permute interaction calculations by distributing
  randomised data-sets across the cores of a computer cluster.
</p>

<h2>Mode Parameter</h2>
<p>The invocation of this epistasis program is controlled by the mode parameter which controls
what the program does and how many cores the program should run on.
The values of the mode parameter and the required number of cores is listed below.
</p>

<table class=options border=0>
<tr><th class=param_header colspan=1>Value</th>
<th class=param_header colspan=1>Description</th>
<th class=param_header colspan=1>No. Cores</th>
<th class=param_header colspan=1>End-point</th></tr>

<tr>
<td class=param_desc>1</th>
<td class=param_desc>Reads a SNP data-set and distributes the interaction calculations across 
the cores of a cluster.</td>
<td class=param_desc>Multi-Core</td>
<td class=param_desc>No</td>
</tr>

<tr>
<td class=param_desc>2</th>
<td class=param_desc>
Merges the interaction data into a single result-set,
calculates the adjusted p-values for all interactions
and applies threshold to find interesting interactions in a data-set.
</td>
<td class=param_desc>Single-Core</td>
<td class=param_desc>No</td>
</tr>

<tr>
<td class=param_desc>3</th>
<td class=param_desc>
Write out the interesting interactions of the non-permuted calculation
to a flat file.
</td>
<td class=param_desc>Single-Core</td>
<td class=param_desc>Yes</td>
</tr>

<tr>
<td class=param_desc>4</th>
<td class=param_desc>
Takes interesting interactions found by the non-permuted
and perform whole data-set permutations across multiple cores of a cluster.
</td>
<td class=param_desc>Multi-Core</td>
<td class=param_desc>No</td>
</tr>

<tr>
<td class=param_desc>5</th>
<td class=param_desc>
Merge the permutation results and writes output to flat file.
</td>
<td class=param_desc>Single-Core</td>
<td class=param_desc>Yes</td>
</tr>
</table>

<h2>OPTIONS</H2>

<table class=options border=0>
<tr><th class=param_header colspan=1>PARAMETER</th>
<th class=param_header colspan=3>DESCRIPTION</th></tr>

<tr><td class=param_name colspan=1>--id STRING</td>
<td class=param_desc colspan=3>Unique identifer for a cluster job (also sets output file stem).</td></tr>

<tr><td class=param_name colspan=1>--rdir PATH</td>
<td class=param_desc colspan=3>Path to remote working directory on central file system.</td></tr>

<tr><td class=param_name colspan=1>--ldir PATH</td>
<td class=param_desc colspan=3>Path to local working directory on each core of the cluster.</td></tr>

<tr><td class=param_name colspan=1>--bmap FILE</td>
<td class=param_desc colspan=3>Path to the BMAP file (binary-encoded SNP list).</td></tr>

<tr><td class=param_name colspan=1>--samples FILE</td>
<td class=param_desc colspan=3>Path to the Sample file (people and phenotype).</td></tr>

<tr><td class=param_name colspan=1>--sped FILE</td>
<td class=param_desc colspan=3>Path to the SPED file (Allelic data, SNP Major, bit-strings).</td></tr>

<tr><td class=param_name colspan=1>--nperm 1..10000</td>
<td class=param_desc colspan=3>
Specify the number of permutations to run on the cluster.
Specifying permutation number causes this programs to run randomised data-sets 
on each core of the cluster.
</td></tr>

<tr><td class=param_name colspan=1>--missed-rank 0..101</td>
<td class=param_desc colspan=3>%Set processor rank to number specified by this parameter instead of rank assigned by MPI.<br>
To be used only if a section of a calculation needs to run a single core if a processor failed.
</td></tr>

<tr><td class=param_name colspan=1>--missed-procs 1..102</td>
<td class=param_desc colspan=3>%Set the No. processors associated with a calculation incase processor rank set by <em>missed-rank</em>.</td></tr>

<tr><td class=param_name colspan=1>--mode 1..5</td>
<td class=param_desc colspan=3>Invocation mode for the epistasis mode</td></tr>
</table>
*/

#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include "SMG_All.h"

int fast_pepi_v2_main(int argc, char *argv[]);
void fast_pepi_v2_write_header(FILE *out);
void fast_pepi_v2_write_record(FILE *out, char *namp1, char *namp2, struct adjusted_2d *adj, double emp1, double emp2);
void fast_pepi_v2_write_header_with_position(FILE *out);
void fast_pepi_v2_write_record_with_position(FILE *out, struct small_locus *snp1, struct small_locus *snp2, struct adjusted_2d *adj, double emp1, double emp2);
int main(int argc, char *argv[]);

struct option long_options[] = {
	{ "bmap", required_argument, NULL, 'm' }, // Path to the map file.
	{ "samples", required_argument, NULL, 's' }, // Path to the map file.
	{ "sped", required_argument, NULL, 'p' }, // Path to the PED file.
	{ "id", required_argument, NULL, 'i' }, // Run ID
	{ "rdir", required_argument, NULL, 'r'}, // Remote Output direct
	{ "ldir", required_argument, NULL, 'l'}, // Local Output direct
	{ "missed-rank", required_argument, NULL, 'n'}, // A rank that failed for whatever reason that needs re-running.
	{ "missed-procs", required_argument, NULL, 'o'}, // No. processors associated with the original calculation.
	{ "nperm", required_argument, NULL, 'P'}, // No. permutations.
	{ "epi_alpha1", required_argument, NULL, 'e' }, // Tight EPI threshold
	{ "mode", required_argument, NULL, 'M'},
	{ "sort-by", required_argument, NULL, 'S' }, // Column Name to sort final output by.
	{ 0,    0,    0,    0   } 
};

/*!
\brief Main body of the FAST_PEPI_BT_V2 program
\ingroup fast_pepi_bt_v2
*/
int fast_pepi_v2_main(int argc, char *argv[]) {
	char ch;
	struct selected_options ops;
	int myrank, numprocs, result;
	struct chromosomes chromos;
	BOOL jobDone = FALSE, *is_valids = NULL, is_valid = FALSE;
	struct small_locus *map = NULL;
	struct sample *samples = NULL;
	char *spedData = NULL;
	int nMap = 0, nSamples = 0, indices[2], missedRank = -1, missedProcs = -1;
	int spedNSamples = 0, nSNPs = 0, nSpedRecordSizeInBytes = 0, nInterestingInteractions = 0, p1;
	char localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH];
	register int p2;
	int *seeds = NULL, nSeeds = 0, *perms = NULL, mode = MODE_NOT_SET, nRaw_data = 0, nAdj = 0, *R = NULL, *maxR = NULL, *nInts = NULL;
	struct small_int2 *raw_data = NULL;
	struct adjusted_2d *adj = NULL;

	// Initialise MPI environment, get size and rank.
	if ((result = MPI_Init(&argc,&argv)) != MPI_SUCCESS) mpi_error(result, "Initialisation");
	if ((result = MPI_Comm_size(MPI_COMM_WORLD,&numprocs)) != MPI_SUCCESS) mpi_error(result, "Number of processors");
	if ((result = MPI_Comm_rank(MPI_COMM_WORLD,&myrank)) != MPI_SUCCESS) mpi_error(result, "Processor number");

	options_init(&ops);
	chromosomes_init(&chromos);
	chromosomes_define_human(&chromos);
	init_genrand((unsigned long) time(NULL));
	*ops.ldir = *ops.szRunId = '\0'; 
	ops.verbose = FALSE;

	while ((ch = (char) getopt_long(argc, argv, "m:v:p:f:i:", long_options, NULL)) != -1) {
		double value;

		switch (ch) {
			case 'i':
				strcpy(ops.szRunId, optarg);
				break;
			case 's': // Sample File
				strcpy(ops.szSampleFileName, optarg);
				break;
			case 'p': // sped file
				strcpy(ops.szSpedFileName, optarg);
				break;
			case 'm': // Bmap File
				strcpy(ops.szBmapFileName, optarg);
				break;
			case 'r': // Remote output directory
				strcpy(ops.rdir, optarg);
				break;
			case 'l': // Local output directory
				strcpy(ops.ldir, optarg);
				break;
			case 'n': // Missed rank to re-run calculation on a single processor.
				missedRank = atoi(optarg);
				break;
			case 'o': // No. processors associated with a 'missed' rank.
				missedProcs = atoi(optarg);
				break;
			case 'P': // Number of Permutations.
				ops.nperm = atoi(optarg);
				if (ops.nperm <= 0) ops.nperm = DEFAULT_NPERM;
				break;
			case 'e': // Tight EPI threshold (epi_alpha1).
				value = atof(optarg);
				if (!(value > 0 && value <= 0.95)) value = DEFAULT_EPI_ALPHA1;
				ops.epi_alpha1 = value;
				break;
			case 'M': // Run-Mode
				mode = atoi(optarg);
				break;
			case 'S': // Sort column name
				if (strlen(optarg) > 0) strcpy(ops.szSortColumnName, optarg);
				break;
			default:
				break;
		}
	}

	if (mode == MODE_NOT_SET) {
		chromosomes_free_attributes(&chromos);
		MPI_Finalize();
		return -1;
	}

	// See if a calculation needs to be re-run on a 'single' processor.
	// Check that all the command-line options are set.
	if (missedRank >= 0 || missedProcs >= 0) {
		BOOL exitEarly = FALSE;

		if (numprocs != 1) exitEarly = TRUE;
		else if (missedRank == -1 || missedProcs == -1) exitEarly = TRUE;
		else if (!(missedRank >= 0 && missedRank <= missedProcs)) exitEarly = TRUE;

		if (exitEarly) {
			chromosomes_free_attributes(&chromos);
			MPI_Finalize();
			return -1;
		} else {
			myrank = missedRank;
			numprocs = missedProcs;
		}
	} 

	// The non-permuted calculation spread out across 'N' ranks.
	if (mode == MODE_NON_PERMUTED) {
		if ((map = bmapfile_load(&ops, &nMap)) != NULL && (samples = samplefile_load(&ops, &nSamples)) != NULL &&
			(spedData = spedfile_load(&ops, &nSNPs, &spedNSamples, &nSpedRecordSizeInBytes)) != NULL) 
		{
			// Allocate array for the entire interaction data-set, this can get quite big.
			raw_data = small_int2_create_array((int) getPairwiseTests(nMap), TRUE);
			if (raw_data != NULL) {
				size_t offset = nSpedRecordSizeInBytes * sizeof(char);
				int int2_index = 0;

				// Array sizes must match so nothing weird is going on.
				if (nMap == nSNPs && spedNSamples == nSamples && epi_util_set_2d_limits(myrank, nSNPs, numprocs, indices)) {
					int interaction_index = 0;
					BOOL doOuter = TRUE;

					// Binary trait, so apply affected status on sample array.
					samplefile_aff_coding(samples, nSamples);
					samplefile_set_original_order(samples, nSamples);

					// Use pointer arithmetic to get the SNP data out of the binary input file.
					for (p1 = 0; p1 < nMap && doOuter; p1++) {
						BitString bs_snp1 = (BitString) (spedData + (p1 * offset)); // SNP1 data

						for (p2 = 0; p2 < nSNPs; p2++) {
							BitString bs_snp2 = (BitString) (spedData + (p2 * offset)); // SNP2 data
							double z;
							
							if (p1 >= p2) continue;
							else if (p1 == p2) continue;

							// Only do the alloted 2-D array slice.
							if (interaction_index < indices[0]) {
								interaction_index++;
								continue;
							} else if (interaction_index > indices[1]) {
								doOuter = FALSE;
								break;
							}

							// Skip X chromosome for now
							if (chromos.chr_sex[map[p1].chr] || chromos.chr_sex[map[p2].chr] || 
								chromos.chr_haploid[map[p1].chr] || chromos.chr_haploid[map[p2].chr]) 
								continue;

							z = fast_pepi_assoc(p1, bs_snp1, p2, bs_snp2, samples, nSamples, NULL, NULL, NULL, NULL, -1, &is_valid);
							if (is_valid) {
								raw_data[int2_index].p1 = p1;
								raw_data[int2_index].p2 = p2;
								raw_data[int2_index].z = z;
								int2_index++;
							}

							interaction_index++;
						}
					}
			
					if (int2_index == 0) jobDone = TRUE; // no valid interactions
					else if (int2_index > 0) {
						// Write out the section of raw epistasis data to file.
						if (small_int2_write_and_copy(numprocs, myrank, int2_index, raw_data, &ops)) jobDone = TRUE;
					}
				}
			}
		} 
	} else if (mode == MODE_NON_PERMUTED_MERGE && numprocs == 1) {
		// Need to read the map first to get the expected data-set size.
		if ((map = bmapfile_load(&ops, &nMap)) != NULL) {
			if ((raw_data = small_int2_load(nMap, &ops, &nRaw_data, &numprocs)) != NULL) {
				qsort(raw_data, nRaw_data, sizeof(struct small_int2), small_int2_cmp); 
				// Scan the result-set looking for interesting values.
				is_valids = create_bool_array(nRaw_data, FALSE);
				if (is_valids != NULL) {
					ops.epi_alpha1 = fabs(ltqnorm(ops.epi_alpha1 / 2));
					for (p1 = 0; p1 < nRaw_data; p1++) {
						if (raw_data[p1].z >= ops.epi_alpha1) {
							is_valids[p1] = TRUE;
							nInterestingInteractions++;
						}
					}
					
					// Calculate the adjusted p-values but only if there are interesting interactions.
					// If there isn't then there's not much point doing it then.
					if (nInterestingInteractions > 0) {
						adj = adjusted_2d_create_array(nRaw_data);
						if (adj != NULL) {
							for (p1 = 0; p1 < nRaw_data; p1++) {
								adj[p1].p1 = raw_data[p1].p1;
								adj[p1].p2 = raw_data[p1].p2;
								if (raw_data[p1].z >= 0) {
									adj[p1].z = raw_data[p1].z;
									adj[p1].p = normdist(-adj[p1].z) * 2;
								}
							}
							if (adjusted_2d_calc(adj, nRaw_data)) {
								if (adjusted_2d_write_and_copy(adj, nRaw_data, is_valids, &ops, myrank)) {
									epi_utils_cleanup(&ops, mode, numprocs, TRUE);
									jobDone = TRUE;
								}
							}
						}
					} 
				}
			}
		}
	} else if (mode == MODE_STOP && numprocs == 1) {
		map = bmapfile_load(&ops, &nMap);
		adj = adjusted_2d_load(&ops, &nAdj);

		// Write out the epistasis results for the non-permuted calculation.
		if (map != NULL && adj != NULL && nMap > 0 && nAdj > 0) {
			FILE *out = NULL;
			sprintf(localFilepath, "%s%s%s_%i", ops.ldir, ops.szRunId, SUFFIX_FEPI_FILE, myrank);
			sprintf(remoteFilepath, "%s%s%s", ops.rdir, ops.szRunId, SUFFIX_FEPI_FILE);
			if ((out = fopen(localFilepath, "w")) != NULL) {
				fast_pepi_v2_write_header_with_position(out);
				for (p1 = 0; p1 < nAdj; p1++) fast_pepi_v2_write_record_with_position(out, &map[adj[p1].p1], &map[adj[p1].p2], &adj[p1], 1.0, 1.0);
				fclose(out);
				if (sort_file_by_col_and_copy(localFilepath, remoteFilepath, ops.szSortColumnName)) {
					epi_utils_cleanup(&ops, mode, numprocs, TRUE);
					jobDone = TRUE;
				}
			}
		} 
	} else if (mode == MODE_PERMUTED) {
		// Do Epistasis on permuted data-sets.
		if ((map = bmapfile_load(&ops, &nMap)) != NULL && 
			(samples = samplefile_load(&ops, &nSamples)) != NULL &&
			(spedData = spedfile_load(&ops, &nSNPs, &spedNSamples, &nSpedRecordSizeInBytes)) != NULL &&
			(seeds = seedfile_load(&ops, &nSeeds)) != NULL &&
			(adj = adjusted_2d_load(&ops, &nAdj)) != NULL) 
		{
			long offset = nSpedRecordSizeInBytes * sizeof(char);
			int nPermsPerProcessor = get_tasks_per_processor(numprocs, ops.nperm);
			int start_perm, end_perm, i;

			// Calculate indexes for the outer permutation loop.
			start_perm = myrank * nPermsPerProcessor;
			end_perm = start_perm + nPermsPerProcessor;
			if (end_perm > ops.nperm || myrank == numprocs - 1) end_perm = ops.nperm;
			perms = create_int_array(nSamples, 0); // Permutation indices.
			R = create_int_array(nAdj, 0);
			maxR = create_int_array(nAdj, 0);
			is_valids = create_bool_array(nAdj, FALSE);
			nInts = create_int_array(nAdj, 0);
			
			if (perms != NULL && nMap == nSNPs && nSeeds >= numprocs && spedNSamples == nSamples && R != NULL && maxR != NULL && is_valids != NULL) {				
				qsort(adj, nAdj, sizeof(struct adjusted_2d), adjusted_2d_cmp);
				samplefile_set_original_order(samples, nSamples);
				samplefile_aff_coding(samples, nSamples);
				init_genrand(seeds[myrank]);

				// Do calculation on shuffled data-set.
				for (i = start_perm; i < end_perm; i++) {
					double zMaxForWholeDataSet = -1000;
					init_bool_array(is_valids, nAdj, FALSE);

#ifdef ECHO_PERMUTATION
					if (myrank == ROOT) {
						fprintf(stdout, "Permutation: %i\r", i + 1);
						fflush(stdout);
					}
#endif
					// Randomise the sample phenotype.
					permute(perms, nSamples);
					samplefile_permute_phenotype(samples, perms, nSamples);

					for (p1 = 0; p1 < nSNPs; p1++) {
						BitString bs_snp1 = (BitString) (spedData + (p1 * offset)); // SNP1 data

						for (p2 = 0; p2 < nSNPs; p2++) {
							BitString bs_snp2;
							double z;

							if (p1 >= p2) continue;
							else if (p1 == p2) continue;

							// Skip X chromosome for now
							if (chromos.chr_sex[map[p1].chr] || chromos.chr_sex[map[p2].chr] || 
								chromos.chr_haploid[map[p1].chr] || chromos.chr_haploid[map[p2].chr]) 
								continue;

							bs_snp2 = (BitString) (spedData + (p2 * offset)); // SNP2 data
							z = fast_pepi_assoc(p1, bs_snp1, p2, bs_snp2, samples, nSamples, NULL, NULL, NULL, NULL, -1, &is_valid);
							if (is_valid) {
								if (z > zMaxForWholeDataSet) zMaxForWholeDataSet = z;
								
								// Update the 'interesting interactions' result array.
								if (z >= ops.epi_alpha && adj != NULL) {
									int index = -1;
									struct adjusted_2d interaction;

									interaction.p1 = p1;
									interaction.p2 = p2;
									index = adjusted_2d_search(adj, nAdj, &interaction); 
									
									if (index >= 0) {
										is_valids[index] = TRUE;
										nInts[index]++;
										if (z > adj[index].z) R[index]++; // EMP1 counts
									}
								}
							}
						}
					}

					// EMP2 counts (FWER)
					for (p1 = 0; p1 < nAdj; p1++) if (is_valids[p1]) if (zMaxForWholeDataSet > adj[p1].z) maxR[p1]++;
				}
				if (epi_util_perm_counts_write_and_copy(&ops, numprocs, start_perm, end_perm, adj, nAdj, R, maxR, nInts, myrank)) jobDone = TRUE;
			}
		}
	} else if (mode == MODE_PERMUTED_MERGE && numprocs == 1) {
		map = bmapfile_load(&ops, &nMap);
		adj = adjusted_2d_load(&ops, &nAdj);
		if (map != NULL && adj != NULL && nMap > 0 && nAdj > 0) {
			// Permutation count arrays.
			R = create_int_array(nAdj, 0);
			maxR = create_int_array(nAdj, 0);
			nInts = create_int_array(nAdj, 0);

			if (epi_util_perm_counts_read(&ops, adj, nAdj, R, maxR, nInts, &numprocs)) {
				FILE *out = NULL;
				sprintf(localFilepath, "%s%s%s", ops.ldir, ops.szRunId, SUFFIX_FEPI_FILE);
				sprintf(remoteFilepath, "%s%s%s", ops.rdir, ops.szRunId, SUFFIX_FEPI_FILE);
				out = fopen(localFilepath, "w");
				if (out != NULL) {
					fast_pepi_v2_write_header_with_position(out);
					for (p1 = 0; p1 < nAdj; p1++) {
						if (nInts[p1] > 0) {
							double emp1 = (double) (R[p1] + 1) / (double) (nInts[p1] + 1);
							double emp2 = (double) (maxR[p1] + 1) / (double) (nInts[p1] + 1);
							fast_pepi_v2_write_record_with_position(out, &map[adj[p1].p1], &map[adj[p1].p2], &adj[p1], emp1, emp2);
						}
					}
					fclose(out);
					jobDone = sort_file_by_col_and_copy(localFilepath, remoteFilepath, ops.szSortColumnName) ? TRUE : FALSE;
					if (jobDone) epi_utils_cleanup(&ops, mode, numprocs, TRUE);
				}
			}
		}
	}
	
	if (samples != NULL) free(samples);
	if (spedData != NULL) free(spedData);
	if (map != NULL) free(map); 
	if (seeds != NULL) free(seeds);
	if (perms != NULL) free(perms);
	if (is_valids != NULL) free(is_valids);
	if (raw_data != NULL) free(raw_data);
	if (adj != NULL) free(adj);
	if (R != NULL) free(R);
	if (maxR != NULL) free(maxR);
	if (nInts != NULL) free(nInts);
	chromosomes_free_attributes(&chromos); 

	MPI_Finalize();

	if (jobDone) return 0;
	else return -1;
}

void fast_pepi_v2_write_header(FILE *out) {
	if (out != NULL) fprintf(out, "SNP1,SNP2,CHISQ,P,EMP1,EMP2,BONF,HOLM,FDR_BH\n");
}

void fast_pepi_v2_write_record(FILE *out, char *name1, char *name2, struct adjusted_2d *adj, double emp1, double emp2) {
	if (out != NULL && adj != NULL) {
		fprintf(out, "%s,%s,%10.16f,%1.16f,%1.16f,%1.16f,%1.16f,%1.16f,%1.16f\n", name1, name2, adj->z * adj->z, adj->p, emp1, emp2, adj->bonf, adj->holm, adj->fdr_bh);
	}
}

void fast_pepi_v2_write_header_with_position(FILE *out) {
	if (out != NULL) fprintf(out, "CHR_SNP1,CHR_SNP2,BP_SNP1,BP_SNP2,SNP1,SNP2,CHISQ,P,EMP1,EMP2,BONF,HOLM,FDR_BH\n");
}

void fast_pepi_v2_write_record_with_position(FILE *out, struct small_locus *snp1, struct small_locus *snp2, struct adjusted_2d *adj, double emp1, double emp2) {
	if (out != NULL && adj != NULL) {
		fprintf(out, "%i,%i,%i,%i,%s,%s,%10.16f,%1.16f,%1.16f,%1.16f,%1.16f,%1.16f,%1.16f\n", snp1->chr, snp2->chr, snp1->bp, snp2->bp, snp1->name, snp2->name, adj->z * adj->z, adj->p, emp1, emp2, adj->bonf, adj->holm, adj->fdr_bh);
	}
}


int main(int argc, char *argv[]) {
	return fast_pepi_v2_main(argc, argv);
}